"""
@Time: 2021/1/14 下午 8:18
@Author: jinzhuan
@File: bert_ee_cn.py
@Desc: 
"""
import sys

sys.path.append('/data/zhuoran/code/cognlp')
sys.path.append('/data/zhuoran/cognlp')

from cognlp import *
import torch
import torch.nn as nn
import torch.optim as optim
from cognlp.core.metrics import EventMetric
from cognlp.core.trainer import Trainer
from cognlp.io.loader.ee.ace2005 import ACE2005Loader
from cognlp.io.processor.ee.ace2005_cn import ACE2005CNProcessor
import numpy as np
from torch.utils.data import RandomSampler


def get_samples_weight(datable, weight=5.0):
    samples_weight = []
    for triggers in datable['triggers']:
        not_none = False
        for trigger in triggers:
            if trigger != 'O':
                not_none = True
                break
        if not_none:
            samples_weight.append(weight)
        else:
            samples_weight.append(1.0)
    return np.array(samples_weight)


torch.cuda.set_device(6)
device = torch.device('cuda')

loader = ACE2005Loader()
train_data, dev_data, test_data = loader.load_all('../../../cognlp/data/ee/ace2005-cn/data')
processor = ACE2005CNProcessor(trigger_path='../../../cognlp/data/ee/ace2005-cn/data/trigger_vocabulary.txt',
                               argument_path='../../../cognlp/data/ee/ace2005-cn/data/argument_vocabulary.txt',
                               max_length=256)
trigger_vocabulary = Vocabulary.load('../../../cognlp/data/ee/ace2005-cn/data/trigger_vocabulary.txt')
argument_vocabulary = Vocabulary.load('../../../cognlp/data/ee/ace2005-cn/data/argument_vocabulary.txt')

train_datable = processor.process(train_data)
train_dataset = DataTableSet(train_datable, to_device=False)
train_sampler = RandomSampler(train_dataset)

dev_datable = processor.process(dev_data)
dev_dataset = DataTableSet(dev_datable, to_device=False)

test_datable = processor.process(test_data)
test_dataset = DataTableSet(test_datable, to_device=False)

model = Bert4EE(trigger_vocabulary=trigger_vocabulary,
                argument_vocabulary=argument_vocabulary, bert_model="hfl/chinese-roberta-wwm-ext")
loss = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(model.parameters(), lr=0.000005)
metric = EventMetric(trigger_vocabulary, argument_vocabulary)

trainer = Trainer(model,
                  train_dataset,
                  dev_data=test_dataset,
                  n_epochs=50,
                  batch_size=10,
                  loss=loss,
                  optimizer=optimizer,
                  scheduler=None,
                  metrics=metric,
                  train_sampler=train_sampler,
                  dev_sampler=None,
                  drop_last=False,
                  gradient_accumulation_steps=1,
                  num_workers=None,
                  save_path='../../../cognlp/data/ee/ace2005-cn/model',
                  save_file=None,
                  print_every=None,
                  scheduler_steps=None,
                  validate_steps=None,
                  save_steps=None,
                  grad_norm=15.0,
                  use_tqdm=True,
                  device=device,
                  device_ids=[6],
                  collate_fn=train_dataset.to_dict,
                  callbacks=None,
                  metric_key=None,
                  writer_path='../../../cognlp/data/ee/ace2005-cn/tensorboard',
                  fp16=False,
                  fp16_opt_level='O1',
                  checkpoint_path='../../../cognlp/data/ee/ace2005-cn/model/ace2005-event-chinese/2021-01-26-21:41:14/checkpoint-1562',
                  task='ace2005-event-chinese',
                  logger_path='../../../cognlp/data/ee/ace2005-cn/logger')

trainer.train()
